# Import all the necessary libraries
from pathlib import Path
from typing import Union

import numpy as np
import pandas as pd
import seaborn as sns
import statsmodels.api as sm
import matplotlib.pyplot as plt

%matplotlib inline
# should enable plotting without explicit call .show()

# For html page generation.
%load_ext pretty_jupyter

def parse_data(data_path: Union[str, Path]) -> pd.DataFrame:
    """
    Parses the data from the given participant, sets new column with participant 
    ID and extracts only the necessary data from the dataset.
    """
    
    # Get participant ID
    participant = data_path.stem.split("_")[1]
    
    # Read the load the csv file.
    df = pd.read_csv(data_path)
    
    # Get only rows with the response data.
    df = df[df.trial_type == "canvas-keyboard-response"]
    
    # Rename the columns for the future usage.
    df = df.rename(
            {
                "rt": "response_time",
                "trial_index": "trial",
                "present": "target_present",
                "setSize": "size",
                "conjuction": "conjuction",
                "correct": "correct",
            },
            axis = 1
        )
    
    # Get only usefull columns.
    df = df\
        [
            [
                "response_time",
                "trial",
                "target_present",
                "size",
                "conjuction",
                "correct",
            ]
        ]
    
    # Add participant ID to data.
    orig_cols = df.columns.to_list()
    new_cols = {
        "participant": participant,
    }
    df = df.assign(**new_cols)
    df = df[["participant"] + orig_cols]

    # Define parameters types.
    df.astype(
        {
            "participant": str,
            "response_time": int,
            "trial": int,
            "target_present": bool,
            "size": int,
            "conjuction": bool,
            "correct": bool,
        }
    )

    return df.sort_values(by="trial")

def load_data(
        data_dir="data",
    ):
    
    """
    Loads the data from all the participants stored in the shared directory 
    in `.csv` files with filenames `ex_{participant_number}.csv`.
    """
    
    selected_files = (Path.cwd() / data_dir).glob("*.csv")

    return pd.concat((parse_data(f) for f in selected_files))

# Load the data from the participants:
exp_data = load_data()

# Show the first lines of the loaded dataset.
exp_data.head()

# Compute the accuracy of the participants.
accuracy_data = exp_data.groupby("participant")\
    .agg(accuracy=("correct", 'mean'))\
    .round(decimals=3)\
    .reset_index()

# Print the table of accuracies for the participants.
accuracy_data

# Get only trials with correct response.
correct_results = exp_data[exp_data.correct]

# Show first items of the dataset without erroneous trials.
correct_results.head()

def plot_reaction_times(
        data: pd.DataFrame,
        title_str: str="Reaction times of participants by set size",
        x_size=None, # To have the x axis in the same scale for the subsets.
    ):
    """
    Plots reaction times for all participants and number of objects in the picture.
    
    Arguments:
        `data`: Selected data.
        `title_str`: Title of the plot.
        `x_size`: Upper bound of the x axis (`None` - default). Used to maintain the same scale for subsets.
    """
    
    # Sort the data by `participant` and `size` to be in correct order for ploting.
    hlp = data.copy().sort_values(by=["participant", "size"])

    hlp["size"] = hlp["size"].astype(str)

    g = sns.catplot(
        data=hlp, 
        x="response_time",
        y="size",
        hue='participant',
        kind='box',
        aspect=2,
    )
    
    g.ax.xaxis.grid(True)
    g.set_xlabels("Reaction time (ms)")
    g.set_ylabels("Set size")
    g.set(title=title_str);
    
    # If specified -> set right `xlim`.
    if x_size != None:
        g.set(xlim=(0, x_size))

# Plot reaction times for the whole data.
plot_reaction_times(correct_results)

# Plot reaction times for only `conjunction` and `feature` search.
cf_task = lambda x: 'conjunction' if x else 'feature'

for is_conjuction in [True, False]:
    plot_reaction_times(
            correct_results[correct_results['conjuction'] == is_conjuction],
            title_str=f"Reaction times of participants for '{cf_task(is_conjuction)}' search by set size",
            x_size=correct_results["response_time"].max() + 200,
        )

def perform_regresion(
        data: pd.DataFrame,
        conf_interval=.05,
    ):
    """
    Computes the regresion slope.
    
    Returns: regresion_slope, lower_confidence_interval, upper_confidence_interval 
    """
    
    sizes = data['size']
    response_times = data['response_time']
    
    # Regresion slope calculation.
    model = sm.OLS(response_times, sm.add_constant(sizes)).fit()

    # Confidence interval calculation.
    cfi = model.conf_int(alpha=conf_interval)
    
    return model.params[1], cfi[0][1], cfi[1][1]

participants = correct_results.participant.unique()


# Result structure for different search type data. 
regresion_results_conj = pd.DataFrame(
    columns=[
            "participant",
            "conjunction",
            "size_slope",
            "conf_int_lower",
            "conf_int_upper",
        ]
    )

for participant in participants:
    for is_conjunction in [True, False]:
        subset = correct_results[
                (correct_results.participant == participant) &
                (correct_results.conjuction == is_conjunction)
            ]
        
        reg_slope, conf_low, conf_up = perform_regresion(subset)
        tmp = pd.DataFrame(
                [
                    [
                        participant,
                        # Convert the `bool` variables to `int` to get rid of warning.
                        int(is_conjunction),
                        reg_slope,
                        conf_low,
                        conf_up,
                    ]
                ], 
                columns=regresion_results_conj.columns
            )

        regresion_results_conj = pd.concat([regresion_results_conj, tmp], ignore_index=True)

participants = correct_results.participant.unique()



# Result structure for different target presence data.
regresion_results_target = pd.DataFrame(
    columns=[
            "participant",
            "target_present",
            "size_slope",
            "conf_int_lower",
            "conf_int_upper",
        ]
    )

# Iterate through all participants.
for participant in participants:
    for is_target_present in [True, False]:
        subset = correct_results[
                (correct_results.participant == participant) &
                (correct_results.target_present == is_target_present)
            ]
        
        reg_slope, conf_low, conf_up = perform_regresion(subset)
        tmp = pd.DataFrame(
                [
                    [
                        participant,
                        # Convert the `bool` variables to `int` to get rid of warning.
                        int(is_target_present),
                        reg_slope,
                        conf_low,
                        conf_up,
                    ]
                ], 
                columns=regresion_results_target.columns
            )
        
        regresion_results_target = pd.concat([regresion_results_target, tmp], ignore_index=True)
        

# Return proper variables back to `bool`.
regresion_results_conj["conjunction"] = regresion_results_conj["conjunction"].astype(bool)
regresion_results_target["target_present"] = regresion_results_target["target_present"].astype(bool)

# Show regresion slopes for the `conjunction` and `feature` search.
regresion_results_conj

# Show regresion results for the target present or missing.
regresion_results_target

def plot_size_slope(
        data: pd.DataFrame,
        hue_param: str,
    ):
    
    is_conj_labels = {
            True: "conjunction",
            False: "feature",
        }
    
    is_target_labels = {
            True: "present",
            False: "not present",
        }
    
    select_legend = lambda x: is_conj_labels if x == "conjunction" else is_target_labels

    title_label = {
            "conjunction": "for different search types",
            "target_present": "based on the target presence",
        }
    
    legend_title = {
        "conjunction": "Search type",
        "target_present": "Target present",
    }
    
    # Create labels for the plots.
    hlp = data.copy()
    hlp[hue_param] = hlp[hue_param].map(select_legend(hue_param)).astype(str)

    # Rename the trials to "{trial_number} : {task_type}" (for plotting).
    # hlp["trial"] = after["trial"] + " : " + after["task_type"]

    g = sns.pointplot(
            data=hlp, 
            x="participant", 
            y="size_slope", 
            hue=hue_param,
            join=False,
            dodge=0.1,
        )
    
    # Set title and axes labels.
    g.set(
        xlabel='Participant',
        ylabel='Slope',
        title=f"Regresion slopes and its confidence intervals \n{title_label[hue_param]}"
    )
    
    g.legend(title=legend_title[hue_param])
    
    # Get positions of the points to draw the confident intervals.
    x_coords = []
    y_coords = []
    for point_pair in g.collections:
        for x, y in point_pair.get_offsets():
            x_coords.append(x)
            y_coords.append(y)

    # Calculate the confidence intervals as bar size.
    confidence_intervals = [
            [data.size_slope.mean()-data.conf_int_lower.min()],
            [data.conf_int_upper.max()-data.size_slope.mean()]
        ]

    # Add confidence interval as lines into the plot.
    g.errorbar(x_coords, y_coords, 
        yerr=confidence_intervals,
        ecolor='black',
        fmt=' ',
        zorder=-1
    )

plot_size_slope(regresion_results_conj, "conjunction")

plot_size_slope(regresion_results_target, "target_present")